import os, sys, cv2
from base64 import b64decode, b64encode
import numpy as np
from typing import Dict

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 5)))
sys.path.insert(0, parent_path)

from core.task.modules.processors.base_processor import BaseProcessor

class Base64Decode(BaseProcessor):
    def __init__(self, keys, Base64Decode_transpose=False) -> None:
        super().__init__(keys)
        self.transpose = Base64Decode_transpose

    def _set(self, data: Dict):
        if not data.get("Base64Decode_transpose",None) is None:
            self.transpose = data["Base64Decode_transpose"]

    def _call(self, data: Dict) -> Dict:
        img_data = data[self.keys["in"]]
        img_data = cv2.imdecode(np.frombuffer(b64decode(img_data),np.uint8),1)
        if self.transpose:
            img_data = cv2.cvtColor(img_data, cv2.COLOR_BGR2RGB)
        data[self.keys["out"]] = img_data[None,...]
        return data